###################################
###################################
##Mediation -- Country year level##
###################################
###################################
# Set RNG
set.seed(123)

# subset na rm data
med.dat <- na.omit(full.2003.2018[,c("gid","clear_zoon_confirm","lagclear_zoon_confirm","Forest.coverage_bin","ag_crops","fires.cell.count_anom","NDVI_mean_anom","log.nl","logged.pop","month","COWCODE","year")])
#Create regular dummies for the mediate package
unique_month <- unique(med.dat$month)
#run loop 
for (month in unique_month) {
  new_col_name <- paste0("f_month", month)  # Create a column name
  med.dat[[new_col_name]] <- ifelse(med.dat$month == month, 1, 0)  # Assign 1 if matched, else 0
}

#country
unique_COWCODE <- unique(med.dat$COWCODE)
#run loop 
for (COWCODE in unique_COWCODE) {
  new_col_name <- paste0("f_COWCODE", COWCODE)  # Create a column name
  med.dat[[new_col_name]] <- ifelse(med.dat$COWCODE == COWCODE, 1, 0)  # Assign 1 if matched, else 0
}

#run year 
#Create regular dummies for the mediate package
unique_year <- unique(med.dat$year)
#run loop
for (year in unique_year) {
  new_col_name <- paste0("f_year", year)  # Create a column name
  med.dat[[new_col_name]] <- ifelse(med.dat$year == year, 1, 0)  # Assign 1 if matched, else 0
}

#############
##Ag. areas##
#############
# subset only ag areas
med.dat.ag <- subset(med.dat, (ag_crops>0))

#Run base model
med.fit1 <- lm(NDVI_mean_anom ~ fires.cell.count_anom+
                 lagclear_zoon_confirm+log.nl +logged.pop+
                 f_month2+f_month3+f_month4+f_month5+f_month6+f_month7+
                 f_month8+f_month9+f_month10+f_month11+f_month12+
                 f_COWCODE570+f_COWCODE565+f_COWCODE572+f_COWCODE541+f_COWCODE571+
                 f_COWCODE580+f_COWCODE552+f_COWCODE540+f_COWCODE551+f_COWCODE553+f_COWCODE490+
                 f_COWCODE510+f_COWCODE484+f_COWCODE516+f_COWCODE501+f_COWCODE481+f_COWCODE517+
                 f_COWCODE500+f_COWCODE520+f_COWCODE411+f_COWCODE471+f_COWCODE482+f_COWCODE625+
                 f_COWCODE626+f_COWCODE530+f_COWCODE450+f_COWCODE437+f_COWCODE475+f_COWCODE452+
                 f_COWCODE461+f_COWCODE434+f_COWCODE451+f_COWCODE438+f_COWCODE483+f_COWCODE439+
                 f_COWCODE432+f_COWCODE404+f_COWCODE436+f_COWCODE433+f_COWCODE531+
                 f_COWCODE420+f_COWCODE435+f_COWCODE615+f_COWCODE600+f_COWCODE620+f_COWCODE651+ 
                 f_COWCODE616+f_year2004+f_year2005+f_year2006+f_year2007+f_year2008+f_year2009+
                 f_year2010+f_year2011+f_year2012+f_year2013+
                 f_year2014+f_year2015+f_year2016+f_year2017+f_year2018, data=med.dat.ag)
summary(med.fit1)
#CSEs
med.cse1 <- coeftest(med.fit1, vcov. = vcovCL(med.fit1, cluster = med.dat.ag$gid, type = "HC0"))

#Run fit model
out.fit1 <- lm(clear_zoon_confirm ~ NDVI_mean_anom+fires.cell.count_anom+
                 lagclear_zoon_confirm+log.nl + logged.pop+
                 f_month2+f_month3+f_month4+f_month5+f_month6+f_month7+
                 f_month8+f_month9+f_month10+f_month11+f_month12+
                 f_COWCODE570+f_COWCODE565+f_COWCODE572+f_COWCODE541+f_COWCODE571+
                 f_COWCODE580+f_COWCODE552+f_COWCODE540+f_COWCODE551+f_COWCODE553+f_COWCODE490+
                 f_COWCODE510+f_COWCODE484+f_COWCODE516+f_COWCODE501+f_COWCODE481+f_COWCODE517+
                 f_COWCODE500+f_COWCODE520+f_COWCODE411+f_COWCODE471+f_COWCODE482+f_COWCODE625+
                 f_COWCODE626+f_COWCODE530+f_COWCODE450+f_COWCODE437+f_COWCODE475+f_COWCODE452+
                 f_COWCODE461+f_COWCODE434+f_COWCODE451+f_COWCODE438+f_COWCODE483+f_COWCODE439+
                 f_COWCODE432+f_COWCODE404+f_COWCODE436+f_COWCODE433+f_COWCODE531+
                 f_COWCODE420+f_COWCODE435+f_COWCODE615+f_COWCODE600+f_COWCODE620+f_COWCODE651+ 
                 f_COWCODE616+f_year2004+f_year2005+f_year2006+f_year2007+f_year2008+f_year2009+                  
                 f_year2010+f_year2011+f_year2012+f_year2013+                  
                 f_year2014+f_year2015+f_year2016+f_year2017+f_year2018, data=med.dat.ag)
summary(out.fit1)
#CSEs
out.cse1 <- coeftest(out.fit1, vcov. = vcovCL(out.fit1, cluster = med.dat.ag$gid, type = "HC0"))

#Set seed
set.seed(123)
#Create a loop for simulations, accommodating memory limits
sim_results_cy <- list()
for (i in 1:10) {  # Split 1000 sims into 10 runs of 100 each
  sim_results_cy[[i]] <- mediate(med.fit1, out.fit1, 
                                 treat = "fires.cell.count_anom", 
                                 mediator = "NDVI_mean_anom",
                                 sims = 100,  cluster = med.dat.ag$gid)
  gc()  # Garbage collection after each iteration
}

#Extract average p values
pvals_ag <- med_pvals(sim_results_cy)

################
##Forest areas##
################
#Subset only forest data
med.dat.for <- subset(med.dat, (Forest.coverage_bin>0))

#Run base model
med.fit2 <- lm(NDVI_mean_anom ~ fires.cell.count_anom+
                 lagclear_zoon_confirm+log.nl + logged.pop+
                 f_month2+f_month3+f_month4+f_month5+f_month6+f_month7+
                 f_month8+f_month9+f_month10+f_month11+f_month12+f_COWCODE565+f_COWCODE541+
                 f_COWCODE580+f_COWCODE552+f_COWCODE540+f_COWCODE551+f_COWCODE553+f_COWCODE490+
                 f_COWCODE510+f_COWCODE484+f_COWCODE516+f_COWCODE501+f_COWCODE481+
                 f_COWCODE500+f_COWCODE520+f_COWCODE411+f_COWCODE471+f_COWCODE482+f_COWCODE625+
                 f_COWCODE626+f_COWCODE530+f_COWCODE450+f_COWCODE437+f_COWCODE475+f_COWCODE452+
                 f_COWCODE461+f_COWCODE434+f_COWCODE451+f_COWCODE438+f_COWCODE483+
                 f_COWCODE432+f_COWCODE404+f_COWCODE433+f_COWCODE531+
                 f_year2004+f_year2005+f_year2006+f_year2007+f_year2008+f_year2009+                 
                 f_year2010+f_year2011+f_year2012+f_year2013+                  
                 f_year2014+f_year2015+f_year2016+f_year2017+f_year2018, data=med.dat.for)
summary(med.fit2)
#CSEs
med.cse2 <- coeftest(med.fit2, vcov. = vcovCL(med.fit2, cluster = med.dat.for$gid, type = "HC0"))

#Run fit model
out.fit2 <- lm(clear_zoon_confirm ~ NDVI_mean_anom+fires.cell.count_anom+
                 lagclear_zoon_confirm+log.nl + logged.pop+
                 f_month2+f_month3+f_month4+f_month5+f_month6+f_month7+
                 f_month8+f_month9+f_month10+f_month11+f_month12+f_COWCODE565+f_COWCODE541+
                 f_COWCODE580+f_COWCODE552+f_COWCODE540+f_COWCODE551+f_COWCODE553+f_COWCODE490+
                 f_COWCODE510+f_COWCODE484+f_COWCODE516+f_COWCODE501+f_COWCODE481+
                 f_COWCODE500+f_COWCODE520+f_COWCODE411+f_COWCODE471+f_COWCODE482+f_COWCODE625+
                 f_COWCODE626+f_COWCODE530+f_COWCODE450+f_COWCODE437+f_COWCODE475+f_COWCODE452+
                 f_COWCODE461+f_COWCODE434+f_COWCODE451+f_COWCODE438+f_COWCODE483+
                 f_COWCODE432+f_COWCODE404+f_COWCODE433+f_COWCODE531+f_year2004+f_year2005+f_year2006+f_year2007+f_year2008+f_year2009+                 
                 f_year2010+f_year2011+f_year2012+f_year2013+                  
                 f_year2014+f_year2015+f_year2016+f_year2017+f_year2018, data=med.dat.for)
summary(out.fit2)
#CSEs
out.cse2 <- coeftest(out.fit2, vcov. = vcovCL(out.fit2, cluster = med.dat.for$gid, type = "HC0"))

#Set seed
set.seed(123)
#Create a loop for simulations, accommodating memory limits
sim_results_cy2 <- list()
for (i in 1:10) {  # Split 1000 sims into 10 runs of 100 each
  sim_results_cy2[[i]] <- mediate(med.fit2, out.fit2, 
                                  treat = "fires.cell.count_anom", 
                                  mediator = "NDVI_mean_anom",
                                  sims = 100,  cluster = med.dat.for$gid)
  gc()  # Garbage collection after each iteration
}

#Extract average p values
pvals_for <- med_pvals(sim_results_cy2)

######################
### All other areas###
######################
#Subset all other areas 
med.dat.other <- subset(med.dat, (Forest.coverage_bin==0|ag_crops==0))

#Run base model
med.fit3 <- lm(NDVI_mean_anom ~ fires.cell.count_anom+log.nl + logged.pop+
                 lagclear_zoon_confirm+
                 f_month2+f_month3+f_month4+f_month5+f_month6+f_month7+
                 f_month8+f_month9+f_month10+f_month11+f_month12+f_COWCODE570+f_COWCODE565+f_COWCODE572+f_COWCODE541+f_COWCODE571+f_COWCODE522+
                 f_COWCODE580+f_COWCODE552+f_COWCODE540+f_COWCODE551+f_COWCODE553+f_COWCODE490+
                 f_COWCODE510+f_COWCODE484+f_COWCODE516+f_COWCODE501+f_COWCODE481+f_COWCODE517+
                 f_COWCODE500+f_COWCODE520+f_COWCODE411+f_COWCODE471+f_COWCODE482+f_COWCODE625+
                 f_COWCODE626+f_COWCODE530+f_COWCODE450+f_COWCODE437+f_COWCODE475+f_COWCODE452+
                 f_COWCODE461+f_COWCODE434+f_COWCODE451+f_COWCODE438+f_COWCODE483+f_COWCODE439+
                 f_COWCODE432+f_COWCODE404+f_COWCODE436+f_COWCODE433+f_COWCODE531+
                 f_COWCODE420+f_COWCODE435+f_COWCODE615+f_COWCODE600+f_COWCODE620+f_COWCODE651+ 
                 f_COWCODE616+f_year2004+f_year2005+f_year2006+f_year2007+f_year2008+f_year2009+                  
                 f_year2010+f_year2011+f_year2012+f_year2013+                 
                 f_year2014+f_year2015+f_year2016+f_year2017+f_year2018, data=med.dat.other)
summary(med.fit3)
#CSEs
med.cse3 <- coeftest(med.fit3, vcov. = vcovCL(med.fit3, cluster = med.dat.other$gid, type = "HC0"))

#Run fit model
out.fit3 <- lm(clear_zoon_confirm ~ NDVI_mean_anom+fires.cell.count_anom+
                 lagclear_zoon_confirm+log.nl + logged.pop+
                 f_month2+f_month3+f_month4+f_month5+f_month6+f_month7+
                 f_month8+f_month9+f_month10+f_month11+f_month12+f_COWCODE570+f_COWCODE565+f_COWCODE572+f_COWCODE541+f_COWCODE571+f_COWCODE522+
                 f_COWCODE580+f_COWCODE552+f_COWCODE540+f_COWCODE551+f_COWCODE553+f_COWCODE490+
                 f_COWCODE510+f_COWCODE484+f_COWCODE516+f_COWCODE501+f_COWCODE481+f_COWCODE517+
                 f_COWCODE500+f_COWCODE520+f_COWCODE411+f_COWCODE471+f_COWCODE482+f_COWCODE625+
                 f_COWCODE626+f_COWCODE530+f_COWCODE450+f_COWCODE437+f_COWCODE475+f_COWCODE452+
                 f_COWCODE461+f_COWCODE434+f_COWCODE451+f_COWCODE438+f_COWCODE483+f_COWCODE439+
                 f_COWCODE432+f_COWCODE404+f_COWCODE436+f_COWCODE433+f_COWCODE531+
                 f_COWCODE420+f_COWCODE435+f_COWCODE615+f_COWCODE600+f_COWCODE620+f_COWCODE651+ 
                 f_COWCODE616+f_year2004+f_year2005+f_year2006+f_year2007+f_year2008+f_year2009+                  
                 f_year2010+f_year2011+f_year2012+f_year2013+                  
                 f_year2014+f_year2015+f_year2016+f_year2017+f_year2018, data=med.dat.other)
summary(out.fit3)
#CSEs
out.cse3 <- coeftest(out.fit3, vcov. = vcovCL(out.fit3, cluster = med.dat.other$gid, type = "HC0"))

#Set seed
set.seed(123)
#Create a loop for simulations, accommodating memory limits
sim_results_cy3 <- list()
for (i in 1:20) {  # Split 1000 sims into 20 runs of 50 each
  sim_results_cy3[[i]] <- mediate(med.fit3, out.fit3, 
                                  treat = "fires.cell.count_anom", 
                                  mediator = "NDVI_mean_anom",
                                  sims = 50,  cluster = med.dat.other$gid)
  gc()  # Garbage collection after each iteration
}

#Extract average p values
pvals_oth <- med_pvals(sim_results_cy3)

#Combine all average p-values
pvals.cy <- cbind(pvals_ag, pvals_for, pvals_oth)
pvals.cy

##Combine all estimates together
sim_results_list <- list(sim_results_cy, sim_results_cy2, sim_results_cy3)

#Generate plots -- model name "main"
med_plots(model.name="CY", sim_results_list=sim_results_list)


##Export mediation models
# Open a file connection
sink("tableCY.txt")

# --- First: Export full regression tables (to LaTeX or text) ---
stargazer(med.cse1, out.cse1, med.cse2, out.cse2, med.cse3, out.cse3,
          type = "text", 
          star.cutoffs = c(0.05, 0.01), 
          style = "default")

cat("\n\n% --- Goodness-of-fit stats only ---\n\n")

# --- Second: Export only N, R², and Adjusted R² ---
stargazer(med.fit1, out.fit1, med.fit2, out.fit2, med.fit3, out.fit3,
          type = "text", 
          keep.stat = c("n", "rsq", "adj.rsq"),
          omit = ".*")

# Close the file
sink()
